import os
import aiohttp
from dotenv import load_dotenv

load_dotenv()

class LLMClient:
    """
    与 DeepSeek API 交互的客户端。
    """
    def __init__(self):
        # self.api_key = os.getenv("DEEPSEEK_API_KEY")
        self.api_key = os.getenv("DEEPSEEK_API_KEY")
        self.api_url = "https://api.deepseek.com/chat/completions"

    async def generate_code_diff(self, prompt: str) -> str:
        """
        向 DeepSeek API 发送一个提示，并返回生成的代码 diff。
        """
        if not self.api_key:
            raise ValueError("在 .env 文件中未找到 DEEPSEEK_API_KEY")

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        data = {
            "model": "deepseek-coder",
            "messages": [
                {"role": "system", "content": "你是一个有用的助手，只返回指定 diff 格式的代码修改。"},
                {"role": "user", "content": prompt}
            ],
            "temperature": 0.7,
            "max_tokens": 1024,
        }

        async with aiohttp.ClientSession() as session:
            try:
                async with session.post(self.api_url, headers=headers, json=data) as response:
                    response.raise_for_status()
                    result = await response.json()
                    return result["choices"][0]["message"]["content"]
            except aiohttp.ClientError as e:
                print(f"调用 DeepSeek API 时出错: {e}")
                return ""